#!/bin/python
# coding=utf-8
"""
Copyright (C) 2018, winewolf@qq.com All rights reserved.
"""

import random
import socket
import hashlib
import time
import threading
from Crypto.Cipher import AES, DES


remote_server_port = 80
proxy_server_port = 8080
proxy_sock, remote_sock = None, None
auth_pwd = 'hello baby!'
auth_padding = '89)$Lg'
BUFFER_SIZE = 16 * 1024 * 1024
unique_proxy = None


def md5(_src):
    return hashlib.md5(_src).digest()


def aes_enc(key, plain):
    aes = AES.new(md5(key + auth_padding), AES.MODE_CBC, md5(key))
    pad_len = 16 - len(plain) % 16
    plain += chr(pad_len)*pad_len
    cipher = aes.encrypt(plain)
    return cipher


def aes_dec(key, cipher):
    aes = AES.new(md5(key + auth_padding), AES.MODE_CBC, md5(key))
    plain = aes.decrypt(cipher)
    return plain[:-ord(plain[-1])]


def listen():
    global proxy_sock, remote_sock
    try:
        proxy_sock = socket.socket()
        proxy_sock.bind(('', proxy_server_port))
        proxy_sock.listen(1024)

        remote_sock = socket.socket()
        remote_sock.bind(('', remote_server_port))
        remote_sock.listen(1024)
        return True
    except:
        return False


def auth(proxy):
    auth_random = hashlib.md5(str(random.random())).digest()
    print '[auth] send random = %r' % random
    proxy.send(auth_random)
    print '[auth] wait hash ...'
    hash_peer = proxy.recv(BUFFER_SIZE)
    # .. calculate
    hash_input = auth_random + auth_pwd + auth_padding
    hash_mine = hashlib.md5(hash_input).digest()
    auth_result = hash_mine == hash_peer
    if auth_result:
        print '[auth] hash matched'
    else:
        print '[auth] hash mismatch (%r != %r)' % (hash_mine, hash_peer)
    ack = aes_enc(auth_pwd, "OK" if auth_result else "FAILED")
    print '[auth] send result = %r' % ack
    proxy.send(ack)
    return auth_result


def run_proxy_server():
    global unique_proxy
    while True:
        unique_proxy, address = proxy_sock.accept()
        print 'Proxy from', address
        if not auth(unique_proxy):
            print 'Auth failed @', address
        else:
            print 'Auth OK @', address

        time.sleep(3600)


def run_remote_server():
    while True:
        client, address = remote_sock.accept()
        print 'Connect from', address
        if unique_proxy:
            req = client.recv(BUFFER_SIZE)
            print '[client] << %s' % req
            to_proxy = aes_enc(auth_pwd, req)

            print '[proxy] >> %r' % to_proxy
            unique_proxy.send(to_proxy)
            from_proxy = unique_proxy.recv(BUFFER_SIZE)
            print '[proxy] << %r' % from_proxy

            ack = aes_dec(auth_pwd, from_proxy)
            print '[client] >> %s' % ack
            client.send(ack)
        else:
            client.send('No proxy connect found!')
        client.close()


if listen():
    p = threading.Thread(target=run_proxy_server)
    r = threading.Thread(target=run_remote_server)
    p.start()
    r.start()
    p.join()
    r.join()
